#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2021 wanghch <wanghch@wanghch-pc>
#
# Distributed under terms of the MIT license.

"""

"""
import pandas as pd
# from sklearn.model_selection import train_test_split
import tensorflow as tf
import os
import sys
from tensorflow.keras.layers import *

from feature_utils import *



dfiles = os.listdir('out/')
absfiles = ["out/" + f for f in dfiles]
dfs = []
for af in absfiles:
    df = pd.read_csv(af, names = COLUMNS, header = None)
    dfs.append(df)

df = pd.concat(dfs)
print(df.shape)
# dc1, dc2, dc3 = get_cate_dims()

df['dt'] = pd.to_datetime(df.date)

X = df[df.dt < '2018-01-01']
X2 = df[df.dt >= '2018-01-01']

X_train, X_cat_train, Y_train = X.loc[:, FNAMES], X.loc[:, CAT_FNAMES], X.loc[:, 'target']
X_test,X_cat_test, Y_test = X2.loc[:, FNAMES], X2.loc[:, CAT_FNAMES], X2.loc[:, 'target']

ALL_FNAMES = FNAMES + CAT_FNAMES
input_map = {}
tinput_map = {}

Y_train = X.loc[:, 'target']
Y_test = X2.loc[:, 'target']
for n in ALL_FNAMES:
    input_map[n] = X.loc[:, n]
    tinput_map[n] = X2.loc[:, n]


num_features = []
cate_features = []
for nc in FNAMES:
    # num_features.append(input_map[nc])
    num_features.append(tf.keras.Input(1, name = nc))

cate_map = {}
for nc in CAT_FNAMES:
    # num_features.append(input_map[nc])
    cate_map[nc] = tf.keras.Input(1, name = nc)
    cate_features.append(cate_map[nc])
#print(num_features)
#input_fs = pd.concat(num_features, axis = 1)
input_fs = tf.keras.layers.concatenate(num_features, axis = 1)
#print(input_fs.tail(10))
#print(X_train.tail(10))
print(input_fs.shape)
#print(X_train.shape)
#sys.exit(0)



cate_dims = get_cate_dims()
dim_features = []
for name, cdim in zip(CAT_FNAMES, cate_dims):
    cdl = Embedding(cdim, 128)(cate_map[name])
    dim_features.append(cdl)

# emb_f = tf.keras.layers.concatenate(dim_features)

layers = []
last = input_fs
# last = tf.keras.layers.concatenate([input_fs] + dim_features)
for i in [128,128,128, 128]:
    # layers.append(tf.keras.layers.Dense(i, activation='relu'))
    last = tf.keras.layers.Dense(i, activation='relu')(last)

final = tf.keras.layers.Dense(1, activation='sigmoid')(last)
# model = tf.keras.Sequential(layers)
model = tf.keras.Model(inputs = num_features + cate_features, outputs = final)
model.compile(
    loss='binary_crossentropy',
    optimizer='adam',
    metrics=['accuracy', tf.keras.metrics.AUC(), 'mse'])

model_dir = "model"
callbacks = []
log_dir = "train_logs/"
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1)
callbacks.append(tensorboard_callback)

model.fit(input_map, Y_train, epochs = 10, callbacks=callbacks)
model.save(model_dir)
#model.evaluate(X_test.values, Y_test.values)
